#include <src/pcps.h>
#include <src/input.h>
#include <src/hamiltonian.h>
#include <src/blas_lapack.h>
#include <src/hf.h>
#include <src/mpi_interface.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   constructs the one-body reduced density matrix from the orbital coefficient matrix
///
/// \param[in]     n        number of orbitals
/// \param[in]     np       number of particles
/// \param[in]     c        the matrix of orbital coefficients,  size n*n
/// \param[out]    x        on exit, the one body reduced density matrix, size n*n
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void pcps::hf::build_rdm1(const int n,
                                             const int np,
                                             const std::vector<S> & c,
                                             std::vector<S> & x)
{
  for (int p = 0; p < n; p++)
  for (int q = 0; q < n; q++) {
    x[n*p+q] = pcps::zero<S>();
    for (int i = 0; i < np; i++)
      x[n*p+q] += pcps::conj( c[n*i+p] ) * c[n*i+q];
  }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   constructs the alpha fock operator
///
/// \param[in]     n        number of orbitals
/// \param[in]     ra       the one body reduced density matrix for alpha, size n*n
/// \param[in]     rb       the one body reduced density matrix for beta , size n*n
/// \param[in]     oei      the one body integrals, size n*n
/// \param[in]     tei      the two body integrals, which are assumed real, in 1122 order, size n*n*n*n
/// \param[out]    f        on exit, the matrix for the fock operator, size n*n
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void pcps::hf::build_fock(const int n,
                                                  const std::vector<S> & ra,
                                                  const std::vector<S> & rb,
                                                  const std::vector<S> & oei,
                                                  const std::vector<S> & tei,
                                                        std::vector<S> & f)
{

  // contribution from the one-electron part of the Hamiltonian
  pcps::xcopy(n*n, &oei.at(0), 1, &f.at(0), 1);

  // coulomb contribution from alpha
  pcps::xgemm('N', 'N', n*n, 1, n*n, pcps::unity<S>(), &tei.at(0), n*n, &ra.at(0), n*n, pcps::unity<S>(), &f.at(0), n*n);

  // coulomb contribution from beta
  pcps::xgemm('N', 'N', n*n, 1, n*n, pcps::unity<S>(), &tei.at(0), n*n, &rb.at(0), n*n, pcps::unity<S>(), &f.at(0), n*n);

  // exchange contribution
  for (int p = 0; p < n; p++)
  for (int q = 0; q < n; q++)
  for (int r = 0; r < n; r++)
  for (int s = 0; s < n; s++)
    f[n*r+q] -= tei[n*n*n*p+n*n*q+n*r+s] * ra[n*p+s];

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   constructs the alpha fock operator using the cholesky decomposed integrals
///
/// \param[in]     n            number of orbitals
/// \param[in]     ra           the one body reduced density matrix for alpha, size n*n
/// \param[in]     rb           the one body reduced density matrix for beta , size n*n
/// \param[in]     ham          the hamiltonian
/// \param[out]    f            on exit, the matrix for the fock operator, size n*n
/// \param[out]    intermed_0   workspace
/// \param[out]    intermed_1   workspace
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void pcps::hf::build_fock_cholesky(const int n,
                                                      const std::vector<S> & ra,
                                                      const std::vector<S> & rb,
                                                      const pcps::Hamiltonian<S> & ham,
                                                      std::vector<S> & f,
                                                      std::vector<S> & intermed_0,
                                                      std::vector<S> & intermed_1)
{

  // get the number of cholesky vectors and size the work vectors accordingly
  const int nchol = ham.n_cholesky();
  if (intermed_0.size() != n*n*nchol) intermed_0.assign(n*n*nchol, pcps::zero<S>());
  if (intermed_1.size() != n*n*nchol) intermed_1.assign(n*n*nchol, pcps::zero<S>());

  // one-electron contribution
  for (int p = 0; p < n; p++)
  for (int q = 0; q < n; q++)
    f[p*n+q] = ham.get_oei(p,q);

  // coulomb contribution
  pcps::xscal(n*n, pcps::zero<S>(), &intermed_0.at(0), 1);
  pcps::xaxpy(n*n, pcps::unity<S>(), &ra.at(0), 1, &intermed_0.at(0), 1);
  pcps::xaxpy(n*n, pcps::unity<S>(), &rb.at(0), 1, &intermed_0.at(0), 1);
  pcps::xgemm('T', 'N', nchol, 1, n*n, pcps::unity<S>(),
              &ham.cholesky_vecs().at(0), n*n,
              &intermed_0.at(0), n*n,
              pcps::zero<S>(),
              &intermed_1.at(0), nchol);
  pcps::xgemm('N', 'N', n*n, 1, nchol, pcps::unity<S>(),
              &ham.cholesky_vecs().at(0), n*n,
              &intermed_1.at(0), nchol,
              pcps::unity<S>(),
              &f.at(0), n*n);

  // exchange contribution
  pcps::xgemm('N', 'N', n, n*nchol, n, pcps::unity<S>(),
              &ra.at(0), n,
              &ham.cholesky_vecs().at(0), n,
              pcps::zero<S>(),
              &intermed_1.at(0), n);
  for (int u = 0; u < nchol; u++)
  for (int r = 0; r < n; r++)
  for (int p = 0; p < n; p++)
    intermed_0[n*n*u + n*p + r] = intermed_1[n*n*u + n*r + p];
  pcps::xgemm('N', 'T', n, n, n*nchol, -pcps::unity<S>(),
              &ham.cholesky_vecs().at(0), n,
              &intermed_0.at(0), n,
              pcps::unity<S>(),
              &f.at(0), n);

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   diagonalizes the fock matrix and extracts the new orbital coefficients
///
/// \param[in]     n        number of orbitals
/// \param[in]     np       number of particles
/// \param[in]     f        the matrix for the fock operator, size n*n
/// \param[out]    eval     on exit, the eigenvalues of the the fock matrix, size n
/// \param[out]    c        on exit, the matrix of orbital coefficients, size n*n
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void pcps::hf::diagonalize_fock(const int n,
                                                        const int np,
                                                        const std::vector<S> & f,
                                                        std::vector<double> & eval,
                                                        std::vector<S> & c)
{

  // preparation for diagonalization
  int info;
  const int lwork = 5*n;
  std::vector<S> work(lwork);
  std::vector<double> rwork(3*n);

  // move the fock matrix into the orbital coefficient matrix in preparation for diagonalization
  pcps::xcopy(n*n, &f.at(0), 1, &c.at(0), 1);

  // take the complex conjugate of the fock matrix.  This way when FORTRAN transposes it the
  // resulting matrix will be equal to the original fock matrix
  pcps::xconj(n*n, &c.at(0), 1);

  // diagonalize the fock matrix, after which the eigenvectors will be in the rows of c
  pcps::xsyev('V', 'U', n, &c.at(0), n, &eval.at(0), &work.at(0), lwork, &rwork.at(0), info);
  if (info != 0)
    throw pcps::Exception("pcps::xsyev failed with error code %i in pcps::hf::diagonalize_fock") % info;

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   constructs the integral arrays used in hartree fock calculations
///
/// \param[in]     ham      object containing the hamiltonian information
/// \param[in]     n        number of orbitals
/// \param[out]    oei_a    on exit, the one body integrals for alpha, size n*n
/// \param[out]    oei_b    on exit, the one body integrals for beta , size n*n
/// \param[out]    tei      on exit, the two body integrals, which are assumed real, in 1122 order, size n*n*n*n
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void pcps::hf::build_ints(const pcps::Hamiltonian<S> & ham,
                                                  const int n,
                                                  std::vector<S> & oei_a,
                                                  std::vector<S> & oei_b,
                                                  std::vector<S> & tei)
{

//  // quantum chemistry hamiltonian
//  if ( ham.is_qc() ) {
//
//    // one electron integrals
//    for (int p = 0; p < n; p++)
//    for (int q = 0; q < n; q++)
//      oei_a[p*n+q] = oei_b[p*n+q] = ham.get_oei(p,q);
//
//    // two electron integrals
//    for (int p = 0; p < n; p++)
//    for (int q = 0; q < n; q++)
//    for (int r = 0; r < n; r++)
//    for (int s = 0; s < n; s++)
//      tei[ n*n*n*p + n*n*q + n*r + s ] = ham.get_tei(p,q,r,s);
//
//  // hamiltonian built from individual operators
//  } else {
//
//    // zero the integrals
//    pcps::xscal(n*n, pcps::zero<S>(), &oei_a.at(0), 1);
//    pcps::xscal(n*n, pcps::zero<S>(), &oei_b.at(0), 1);
//    pcps::xscal(n*n*n*n, pcps::zero<S>(), &tei.at(0), 1);
//
//    // declare an iterator for looping over the hamiltonian's operators
//    typename std::vector<typename pcps::OpBase<S>::const_ptr>::const_iterator op;
//
//    // extract integrals from the hamiltonian's operators
//    for (op = ham.ops().begin(); op != ham.ops().end(); op++) {
//
//      // raising/lowering operators
//      if (typeid(**op) == typeid(pcps::OpRL<S>)) {
//
//        // check that we do not have any ab contributions
//        if ( (**op).sites().at(0) % 2 != (**op).sites().at(1) % 2 )
//          throw pcps::Exception("RL operators with one alpha and one beta site cannot be used in pcps::hf::build_ints");
//
//        // add the operator's contribution to the one electron integrals
//        const int site1 = (**op).sites().at(0)/2;
//        const int site2 = (**op).sites().at(1)/2;
//        if ((**op).sites().at(0) % 2 == 0 )
//          oei_a[ n * site1 + site2 ] += (**op).coeff();
//        else
//          oei_b[ n * site1 + site2 ] += (**op).coeff();
//
//      // double number operators
//      } else if (typeid(**op) == typeid(pcps::OpNN<S>)) {
//
//        // we currently only support on-site two electron terms like in the Hubbard model
//        const int s0 = (**op).sites().at(0);
//        const int s1 = (**op).sites().at(1);
//        if ( ! ( s0 % 2 == 0 && s1 == s0 + 1 ) &&
//             ! ( s0 % 2 == 1 && s1 == s0 - 1 )    )
//          throw pcps::Exception("NN operators must act on one spatial site in pcps::hf::build_ints");
//
//        // add the operator's contribution to the two electron integrals
//        const int s = (**op).sites().at(0) / 2;
//        tei[ n*n*n*s + n*n*s + n*s + s ] += (**op).coeff();
//
//      // other operators are not currently supported
//      } else {
//        throw pcps::Exception("operator type \"%s\" not supported by pcps::hf::build_ints") % pcps::type_lookup(typeid(**op));
//      }
//
//    }
//
//  }

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   computes the alpha contribution to the hartree fock energy
///
/// \param[in]     n        number of orbitals
/// \param[in]     ham      the hamiltonian
/// \param[in]     f        the matrix for the fock operator, size n*n
/// \param[in]     rdm      the one body reduced density matrix, size n*n
///
/// \return the alpha contribution to the hartree fock energy
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> S pcps::hf::alpha_energy(const int n,
                                            const pcps::Hamiltonian<S> & ham,
                                            const std::vector<S> & f,
                                            const std::vector<S> & rdm)
{

  // initialize the energy
  S retval = pcps::zero<S>();

  // compute the energy
  for (int p = 0; p < n; p++)
  for (int q = 0; q < n; q++)
    retval += 0.5 * ( ham.get_oei(p,q) + f[p*n+q] ) * rdm[p*n+q];

  // return the energy
  return retval;

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   print the hartree fock orbitals and orbitals energies
///
/// \param[in]     n        number of orbitals
/// \param[in]     na       number of alpha particles
/// \param[in]     evals    the orbital energy eigenvalues, size n
/// \param[in]     c        the matrix of orbital coefficients,  size n*n
/// \param[in]     str      string for the particle type
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void pcps::hf::print_orbitals(const int n,
                                                      const int na,
                                                      const std::vector<double> & evals,
                                                      const std::vector<S> & c,
                                                      const std::string & str)
{

  if (str.size() > 0)
    std::cout << boost::format("MO coefficients for %s:") % str << std::endl;
  else
    std::cout << boost::format("MO coefficients:") << std::endl;
  std::cout << std::endl;
  for (int i = 0; i < n; i++) {
    std::cout << boost::format("%20.9f  ") % evals[i];
    if (i < na)
      std::cout << boost::format("%4s  ") % "occ";
    else
      std::cout << boost::format("%4s  ") % "vir";
    for (int s = 0; s < n; s++)
      std::cout << boost::format(" %s") % pcps::format_number<S>("%12.5f", c[i*n+s]);
    std::cout << std::endl;
  }
  std::cout << std::endl;

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   computes the hartree fock energy and orbital coefficients
///
/// \param[in]     rhf         whether to perform a restricted or unrestricted calculation
/// \param[in]     n           number of orbitals
/// \param[in]     na          number of alpha particles
/// \param[in]     nb          number of beta  particles
/// \param[in]     ham         object containing the hamiltonian information
/// \param[in]     bias        for each orbital, a bias towards alpha or beta, size n
/// \param[in]     bias_decay  decay rate for the orbital bias
/// \param[in]     thresh      energy convergence threshold
/// \param[in]     maxit       maximum number of self consistent iterations
/// \param[in,out] ca          On entry, the initial guess for the alpha orbital coefficients,  size n*n.
///                            On exit, the optimized matrix of alpha orbital coefficients.
/// \param[in,out] cb          On entry, the initial guess for the beta orbital coefficients,  size n*n.
///                            On exit, the optimized matrix of beta orbital coefficients.
/// \param[out]    energy      On exit, the unrestricted hartree fock energy.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void pcps::hf::run_scf(const bool rhf,
                                          const int n,
                                          const int na,
                                          const int nb,
                                          const pcps::Hamiltonian<S> & ham,
                                          const std::vector<double> & bias,
                                          const double bias_decay,
                                          const double thresh,
                                          const int maxit,
                                          std::vector<S> & ca,
                                          std::vector<S> & cb,
                                          S & energy)
{

  // only perform HF on root process
  if (pcps::mpi::rank() == 0) {

    // print a greeting
    if (rhf)
      std::cout << boost::format("starting restricted hartree fock iterations") << std::endl;
    else
      std::cout << boost::format("starting unrestricted hartree fock iterations") << std::endl;
    std::cout << std::endl;

    //// allocate and initialize the integrals
    //std::vector<S> oei_a(n*n, pcps::zero<S>());
    //std::vector<S> oei_b(n*n, pcps::zero<S>());
    //std::vector<S> tei(n*n*n*n, pcps::zero<S>());
    //pcps::hf::build_ints(ham, n, oei_a, oei_b, tei);

    // check that bias is zero for restricted calculations
    for (int i = 0; rhf && i < bias.size(); i++)
      if ( bias.at(i) != 0.0 )
        throw pcps::Exception("bias must be zero for restricted hartree fock");

    // allocate fock matrices
    std::vector<S> fa(n*n, pcps::zero<S>());
    std::vector<S> fb(n*n, pcps::zero<S>());

    // allocate one body reduced density matrices
    std::vector<S> ra(n*n, pcps::zero<S>());
    std::vector<S> rb(n*n, pcps::zero<S>());

    // allocate orbital energy vectors
    std::vector<double> eval_a(n, 0.0);
    std::vector<double> eval_b(n, 0.0);

    // prepare some workspace vectors
    std::vector<S> intermed_0, intermed_1;

    // initialize the decay coefficient for the bias
    double decay = 1.0;

    // iterate
    bool converged = false;
    for (int iter = 0; iter < maxit; iter++) {

      // for restricted calculations, alpha and beta orbitals are the same
      if (rhf)
        cb = ca;

      // build one body density matrices
      pcps::hf::build_rdm1(n, na, ca, ra);
      if (rhf)
        pcps::xcopy(n*n, &ra.at(0), 1, &rb.at(0), 1);
      else
        pcps::hf::build_rdm1(n, nb, cb, rb);

      // build fock matrices
      pcps::hf::build_fock_cholesky(n, ra, rb, ham, fa, intermed_0, intermed_1);
      if (rhf)
        pcps::xcopy(n*n, &fa.at(0), 1, &fb.at(0), 1);
      else
        pcps::hf::build_fock_cholesky(n, rb, ra, ham, fb, intermed_0, intermed_1);

      // compute energy
      const S last_energy = energy;
      energy = pcps::unity<S>() * ham.e_cor_nuc();
      energy += pcps::hf::alpha_energy(n, ham, fa, ra);
      energy += pcps::hf::alpha_energy(n, ham, fb, rb);

      // print energy
      std::cout << boost::format("%4i   %s") % iter % pcps::format_number<S>("%20.12f", energy) << std::endl;

      // check for convergence
      converged = ( iter > 0 && std::abs( energy - last_energy ) < thresh );
      if (converged) break;

      // add the bias to the fock matrices
      for (int p = 0; p < n; p++) {
        fa[n*p+p] -= pcps::unity<S>() * decay * bias.at(p);
        fb[n*p+p] += pcps::unity<S>() * decay * bias.at(p);
      }

      // diagonalize fock matrices to find the new orbital coefficients
      pcps::hf::diagonalize_fock(n, na, fa, eval_a, ca);
      if (!rhf)
        pcps::hf::diagonalize_fock(n, nb, fb, eval_b, cb);

      // reduce the bias based on the user specified decay rate
      decay *= bias_decay;

    }

    // print results
    std::cout << std::endl;
    if (rhf) {
      if (converged)
        std::cout << boost::format("restricted hartree fock has converged") << std::endl;
      else
        std::cout << boost::format("!!! restricted hartree fock failed to converge !!!") << std::endl;
    } else {
      if (converged)
        std::cout << boost::format("unrestricted hartree fock has converged") << std::endl;
      else
        std::cout << boost::format("!!! unrestricted hartree fock failed to converge !!!") << std::endl;
    }
    std::cout << std::endl;

    // for restricted calculations, alpha and beta orbitals are the same
    if (rhf)
      cb = ca;

    // print the mo coefficients and orbital energies
    pcps::hf::print_orbitals(n, na, eval_a, ca, "alpha");
    if (!rhf)
      pcps::hf::print_orbitals(n, nb, eval_b, cb, "beta");

  }

  // send the results to all processes
  pcps::mpi::bcast(energy);
  pcps::mpi::bcast(ca);
  pcps::mpi::bcast(cb);

}

// explicitly instantiate the function
template void pcps::hf::run_scf(const bool rhf,
                                const int n,
                                const int na,
                                const int nb,
                                const pcps::Hamiltonian<double> & ham,
                                const std::vector<double> & bias,
                                const double bias_decay,
                                const double thresh,
                                const int maxit,
                                std::vector<double> & ca,
                                std::vector<double> & cb,
                                double & energy);
template void pcps::hf::run_scf(const bool rhf,
                                const int n,
                                const int na,
                                const int nb,
                                const pcps::Hamiltonian<std::complex<double> > & ham,
                                const std::vector<double> & bias,
                                const double bias_decay,
                                const double thresh,
                                const int maxit,
                                std::vector<std::complex<double> > & ca,
                                std::vector<std::complex<double> > & cb,
                                std::complex<double> & energy);
